// A mixture of expression divergence
#include "ted.hh"
#include "util.hh"
#include <boost/lexical_cast.hpp>
#include <string>
#include <iostream>
#include <fstream>
#include <vector>
#include <algorithm>
#include <boost/ptr_container/ptr_vector.hpp>

template<typename T>
T get_param(const char* arg, std::string emsg)
{
  T val;
  try
  {
    val = boost::lexical_cast<T>(arg);
  } catch (boost::bad_lexical_cast&)
  {
    err_msg(emsg);
    exit(1);
  }
  return val;
}

template<typename T>
void vec_out(std::ofstream& ofs, const std::vector<T>& vec)
{
  if (vec.size() == 0)
    return;
  ofs << vec.at(0);
  for (size_t j = 1; j < vec.size(); ++j)
  {
    ofs << "\t" << vec.at(j);
  }
}

struct sqrt_t
{
  double operator()(const double x) const
  {
    return std::sqrt(x);
  }
};

template<typename T, typename F>
void vec_out(std::ofstream& ofs, const std::vector<T>& vec, F f)
{
  if (vec.size() == 0)
    return;
  ofs << f(vec.at(0));
  for (size_t j = 1; j < vec.size(); ++j)
  {
    ofs << "\t" << f(vec.at(j));
  }
}

////////////////////////////////////////////////////////////////
struct prog_args_t
{

  explicit prog_args_t() :
      x_file(""), y_file(""), output("output"), // default args
      bmin(-2.), bmax(2.),                      // beta in [bmin, bmax]
      repeat(1), tol(1e-3), max_iter(100)       // options
  {
  }

  std::string x_file;
  std::string y_file;
  std::string output;
  double bmin;
  double bmax;
  int repeat;
  double tol;
  int max_iter;
};

void print_help(const char* prog, const prog_args_t& args)
{
  std::cerr << prog << " -x X.txt -y Y.txt -o out" << std::endl;
  std::cerr << "         -bmin bmin -bmax bmax" << std::endl;
  std::cerr << "         -i iter -r repeat -t tol" << std::endl;
  std::cerr << std::endl;
  std::cerr << " X.txt  control expression file (n x p tab-separated)"
      << std::endl;
  std::cerr << " Y.txt  case expression file (n x p tab-separated)"
      << std::endl;
  std::cerr << " out    output files' header (default: " << args.output << ")"
      << std::endl;
  std::cerr << " bmin   min value of beta" << std::endl;
  std::cerr << " bmax   max value of beta" << std::endl;
  std::cerr << std::endl;
  std::cerr << "" << std::endl;
  std::cerr << std::endl;
}

bool parse_args(const int argc, const char* argv[], prog_args_t& args_out)
{
  for (int j = 1; j < argc; ++j)
  {
    std::string curr = argv[j];

    if (curr == "-x" && ++j < argc)
    {
      args_out.x_file = argv[j];
    }
    else if (curr == "-y" && ++j < argc)
    {
      args_out.y_file = argv[j];
    }
    else if (curr == "-o" && ++j < argc)
    {
      args_out.output = argv[j];
    }
    else if (curr == "-i" && ++j < argc)
    {
      args_out.max_iter = get_param<int>(argv[j], "parsing iter");
    }
    else if (curr == "-r" && ++j < argc)
    {
      args_out.repeat = get_param<int>(argv[j], "parsing repeat");
    }
    else if (curr == "-t" && ++j < argc)
    {
      args_out.tol = get_param<double>(argv[j], "parsing tol");
    }
    else if (curr == "-bmin" && ++j < argc)
    {
      args_out.bmin = get_param<double>(argv[j], "parsing beta min");
    }
    else if (curr == "-bmax" && ++j < argc)
    {
      args_out.bmax = get_param<double>(argv[j], "parsing beta max");
    }
  }

  if (args_out.tol <= 0)
    return false;
  if (args_out.x_file.size() == 0 || args_out.y_file.size() == 0)
    return false;
  if (args_out.bmin >= args_out.bmax)
    return false;
  return true;
}

////////////////////////////////////////////////////////////////
int main(const int argc, const char* argv[])
{

  prog_args_t args;
  if (!parse_args(argc, argv, args))
  {
    print_help(argv[0], args);
    exit(1);
  }

  std::string out(args.output);
  const size_t max_iter = args.max_iter;
  const double tol = args.tol;
  const size_t repeat = args.repeat;

  size_t n, p;
  expr_pair_vec_t data_vec;
  if (!read_expr_pair_vec(args.x_file.c_str(), args.y_file.c_str(), data_vec, n,
      p))
  {
    err_msg("failed to read the data");
    return 1;
  }

  const std::string mean_file = out + ".mean";
  const std::string sd_file = out + ".sd";
  const std::string z_file = out + ".z";

  TLOG("data n = " << n << " x p = " << p);

  // Just fit a single model
  expr_pair_set_t pair_set;
  vec2set(data_vec, pair_set);
  func_expr_div_t func(p, args.bmin, args.bmax);
  typedef func_expr_div_t::vec_t real_vec_t;

  for (size_t iter = 1; iter <= max_iter; ++iter)
  {
    func.update(pair_set);
    double diff = empirical_bayes(func);

    TLOG("iter = " << iter << ", penalty static = " << func.get_penalty_static() //
        << ", penalty kinetic = " << func.get_penalty_kinetic());

    if (diff < tol)
      break;
  }

  typedef func_expr_div_t::vec_t real_vec_t;

  const real_vec_t& mean = func.mean_vec();
  const real_vec_t& var = func.var_vec();
  std::ofstream mean_out(mean_file.c_str(), std::ios::out);
  vec_out(mean_out, mean);
  mean_out << std::endl;
  mean_out.close();

  std::ofstream sd_out(sd_file.c_str(), std::ios::out);
  vec_out(sd_out, var, sqrt_t());
  sd_out << std::endl;
  sd_out.close();

  return 0;
}
